from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi_mcp import FastApiMCP
import asyncio
import json
import os

# 确保生成的文档目录存在
# 虽然 MCPDocumentServer 的 init 中也会创建，但在这里明确确保一次更安全
GENERATED_DOCS_DIR = "generated_docs"
if not os.path.exists(GENERATED_DOCS_DIR):
    os.makedirs(GENERATED_DOCS_DIR)

app = FastAPI()

# 创建 MCP 实例
mcp = FastApiMCP(app)

# 假设你已有的文档生成核心逻辑
# 确保 mcp_file.py 存在且 MCPDocumentServer 类可导入
from mcp_file import MCPDocumentServer

doc_server = MCPDocumentServer()

async def sse_event_generator(question, user_id):
    # 1. 模板匹配
    yield f"data: {json.dumps({'stage': 'matching', 'message': '正在匹配模板...'})}\n\n"
    matched_template, match_score = doc_server.template_manager.match_template(question)

    # 2. 内容生成
    yield f"data: {json.dumps({'stage': 'generating', 'message': '正在生成内容...'})}\n\n"
    generated_content = await doc_server.content_generator.generate_content(
        matched_template, question, {"user_id": user_id}
    )

    # 3. 文档生成
    yield f"data: {json.dumps({'stage': 'document', 'message': '正在生成文档...'})}\n\n"
    file_path = doc_server.document_processor.generate_document(
        matched_template, generated_content, user_id
    )

    # 4. 完成
    # 返回相对于静态文件目录的路径，或者完整的可访问 URL
    # 为了简化，这里返回基于静态目录的相对路径，前端需要自己构造完整URL
    accessible_file_path = os.path.join(GENERATED_DOCS_DIR, os.path.basename(file_path))
    yield f"data: {json.dumps({'stage': 'done', 'file_path': accessible_file_path, 'message': '文档生成完成'})}\n\n"

# 使用标准的 FastAPI 装饰器定义路由
# fastapi-mcp 应该会自动发现这个路由并将其作为 MCP 工具
@app.get("/generate_document")
async def generate_document(request: Request, question: str, user_id: str = "user"):
    """
    根据用户问题生成文档并实时推送进度。
    """
    async def event_stream():
        async for event in sse_event_generator(question, user_id):
            yield event
            if await request.is_disconnected():
                break
    return StreamingResponse(event_stream(), media_type="text/event-stream")

# 将生成的文档目录作为静态文件服务
# 用户可以通过 /generated_docs/<文件名> 访问其中的文件
app.mount("/generated_docs", StaticFiles(directory=GENERATED_DOCS_DIR), name="generated_docs")

# 创建 MCP 实例，传入 FastAPI 应用
mcp = FastApiMCP(app)

# 将 MCP 服务器挂载到 FastAPI 应用的 /mcp 路径
mcp.mount()

# 你可以在这里添加其他 FastAPI 路由
# @app.get("/hello")
# async def hello():
#     return {"message": "Hello World"}